## 模板1
import os


def cuda_build_object(target, source, env):
    """
    target : 目标文件
    source : 源文件列表
    """
    # nvcc 编译 .cu 文件
    # 1.编译
    for file in source:
        file = str(file)
        if file.endswith(".cu"):
            file_name, _ = os.path.splitext(file)
            file_obj = f"{file_name}.o"
            build_cmd = f"nvcc -c {file} -o {file_obj} {env['CCFLAGS']}"
            ret = env.Execute(build_cmd)
            if ret != 0:
                raise Exception(f"Failed to compile {file}")


# 1.创建cuda 构建object对象
cuda_builderObject = Builder(action=cuda_build_object, suffix=".out", src_suffix=".cu")
env = Environment()
env.Append(BUILDERS={"cudaObject": cuda_builderObject})

# 2.nvcc 编译 .cu文件 为 .o
env.cudaObject("cuda_fun.o", "cuda_fun.cu")

# 3.编译为程序
env.Program(
    target="main.out",  # 生成的可执行文件名
    source=["cuda_fun.o", "main.cc"],  # 源文件包括main.cc和已有的add.o
    LIBS=["cudart"],  # 链接的库
)
